//==============================================================================
// Copyright (c) 2008-2013 Niall Ryan. All Rights Reserved.
//==============================================================================

#ifndef HYDRA_HASHMAPS_STRIPEDHASHMAP_H
#define HYDRA_HASHMAPS_STRIPEDHASHMAP_H 1

#include "hydra/atomic/Atomic.h"
#include "hydra/hashutil.h"
#include "hydra/locks/TTASLock.h"

namespace Hydra
{
	/**
	 * Hash map where the buckets are protected by a set of striped locks. The number of buckets can grow, but the number of locks is fixed.
	 */
	template<typename K, typename T, typename Hash=HashUtil::HashFunction<K>, typename Pred=HashUtil::EqualTo<K>, typename Lock=TTASLock, unsigned int NumLocks=16>
	class StripedHashMap
	{
	public:
		StripedHashMap(unsigned int numBuckets=16);
		~StripedHashMap();
		bool insert(const K& key, const T& value);
		bool erase(const K& key);
		bool contains(const K& key);
		bool find(const K& key, T* value);
		void setMaxLoadFactor(float maxLoadFactor) { m_maxLoadFactor = maxLoadFactor; }
		float getLoadFactor() const     { return static_cast<float>(m_size) / static_cast<float>(m_numBuckets); }
		unsigned int getSize() const { return m_size; }
		bool isEmpty() const         { return (m_size == 0); }

		static_assert((NumLocks&(NumLocks-1)) == 0, "Number of locks must be a power of 2");

    private:
		void grow();

		struct Pair
		{
			Pair(const K& key, const T& value) : m_key(key), m_value(value) { }
			K m_key;
			T m_value;
		};
		typedef std::vector<Pair> Bucket;

		Atomic<unsigned int> m_size;
		float m_maxLoadFactor;

		Lock m_locks[NumLocks];

		unsigned int m_numBuckets;
		Bucket* m_buckets;
	};

	//==================================================================================================
	//==================================================================================================
	//==================================================================================================

	template<typename K, typename T, typename Hash, typename Pred, typename Lock, unsigned int NumLocks>
	inline StripedHashMap<K, T, Hash, Pred, Lock, NumLocks>::StripedHashMap(unsigned int numBuckets)
	{
		assert((numBuckets&(numBuckets-1)) == 0);
		m_size = 0;
		m_maxLoadFactor = 0.75f;
		m_numBuckets = max(numBuckets, NumLocks); //need to ensure no more than one lock per bucket
		m_buckets = new Bucket[m_numBuckets];
	}

	template<typename K, typename T, typename Hash, typename Pred, typename Lock, unsigned int NumLocks>
	inline StripedHashMap<K, T, Hash, Pred, Lock, NumLocks>::~StripedHashMap()
	{
		delete[] m_buckets;
	}

	template<typename K, typename T, typename Hash, typename Pred, typename Lock, unsigned int NumLocks>
	inline bool StripedHashMap<K, T, Hash, Pred, Lock, NumLocks>::insert(const K& key, const T& value)
	{
		Hash hash;
		unsigned int hashValue = hash(key);
		Lock& lock = m_locks[hashValue & (NumLocks-1)];
		lock.lock();
		Bucket& bucket = m_buckets[hashValue & (m_numBuckets-1)];
		for( Bucket::const_iterator iter=bucket.begin(); iter!=bucket.end(); ++iter )
		{
			Pred pred;
			if( pred(iter->m_key, key) )
			{
				lock.unlock();
				return false;
			}
		}
		bucket.push_back(Pair(key, value));
		lock.unlock();
		++m_size;
		if( getLoadFactor() > m_maxLoadFactor )
			grow();
		return true;
	}

	template<typename K, typename T, typename Hash, typename Pred, typename Lock, unsigned int NumLocks>
	inline bool StripedHashMap<K, T, Hash, Pred, Lock, NumLocks>::erase(const K& key)
	{
		Hash hash;
		unsigned int hashValue = hash(key);
		Lock& lock = m_locks[hashValue & (NumLocks-1)];
		lock.lock();
		Bucket& bucket = m_buckets[hashValue & (m_numBuckets-1)];
		for( Bucket::iterator iter=bucket.begin(); iter!=bucket.end(); ++iter )
		{
			Pred pred;
			if( pred(iter->m_key, key) )
			{
				*iter = bucket.back();
				bucket.pop_back();
				lock.unlock();
				--m_size;
				return true;
			}
		}
		lock.unlock();
		return false;
	}

	template<typename K, typename T, typename Hash, typename Pred, typename Lock, unsigned int NumLocks>
	inline bool StripedHashMap<K, T, Hash, Pred, Lock, NumLocks>::contains(const K& key)
	{
		Hash hash;
		unsigned int hashValue = hash(key);
		Lock& lock = m_locks[hashValue & (NumLocks-1)];
		lock.lock();
		Bucket& bucket = m_buckets[hashValue & (m_numBuckets-1)];
		for( Bucket::const_iterator iter=bucket.begin(); iter!=bucket.end(); ++iter )
		{
			Pred pred;
			if( pred(iter->m_key, key) )
			{
				lock.unlock();
				return true;
			}
		}
		lock.unlock();
		return false;
	}

	template<typename K, typename T, typename Hash, typename Pred, typename Lock, unsigned int NumLocks>
	inline bool StripedHashMap<K, T, Hash, Pred, Lock, NumLocks>::find(const K& key, T* value)
	{
		Hash hash;
		unsigned int hashValue = hash(key);
		Lock& lock = m_locks[hashValue & (NumLocks-1)];
		lock.lock();
		Bucket& bucket = m_buckets[hashValue & (m_numBuckets-1)];
		for( Bucket::const_iterator iter=bucket.begin(); iter!=bucket.end(); ++iter )
		{
			Pred pred;
			if( pred(iter->m_key, key) )
			{
				*value = iter->m_value;
				lock.unlock();
				return true;
			}
		}
		lock.unlock();
		return false;
	}

	template<typename K, typename T, typename Hash, typename Pred, typename Lock, unsigned int NumLocks>
	inline void StripedHashMap<K, T, Hash, Pred, Lock, NumLocks>::grow()
	{
		unsigned int oldNumBuckets = m_numBuckets;

		//acquire all locks in order
		for( unsigned int lockIndex=0; lockIndex<NumLocks; ++lockIndex )
			m_locks[lockIndex].lock();

		//grow unless another thread has beaten us to it
		if( m_numBuckets == oldNumBuckets )
		{
			//create the new set of buckets, twice as many
			Bucket* oldBuckets = m_buckets;
			m_numBuckets *= 2;
			m_buckets = new Bucket[m_numBuckets];

			//rehash the items into the new buckets.
			// possible optimization - power of 2 bucket count means each bucket splits into 2 new buckets, could reduce copying
			Hash hash;
			for( unsigned int oldBucketIndex=0; oldBucketIndex<oldNumBuckets; ++oldBucketIndex )
			{
				Bucket& oldBucket = oldBuckets[oldBucketIndex];
				for( Bucket::const_iterator iter=oldBucket.begin(); iter!=oldBucket.end(); ++iter )
				{
					unsigned int hashValue = hash(iter->m_key);
					Bucket& newBucket = m_buckets[hashValue & (m_numBuckets-1)];
					newBucket.push_back(*iter);
				}
			}

			delete[] oldBuckets;
		}

		//release all locks
		for( unsigned int lockIndex=0; lockIndex<NumLocks; ++lockIndex )
			m_locks[lockIndex].unlock();
	}
}

#endif
